#!/usr/bin/env python3

import argparse
import numpy as np
import pandas as pd
import cmaps
import proplot as plot
from cartopy.io.shapereader import Reader as ShapeReader
from cartopy.feature import ShapelyFeature
from matplotlib.colors import ListedColormap, BoundaryNorm
from datetime import datetime,timedelta

parser = argparse.ArgumentParser(description='Plot EFSO results')
parser.add_argument('-i', '--input', required=True)
parser.add_argument('-o', '--output', required=True)
parser.add_argument('--min-lon', type=float, default=73)
parser.add_argument('--max-lon', type=float, default=135)
parser.add_argument('--min-lat', type=float, default=15)
parser.add_argument('--max-lat', type=float, default=55)
args = parser.parse_args()

proj = plot.Proj('pcarree')

time_i = datetime.strptime(args.input.split('.')[-2],'%Y%m%d%H')

china = ShapelyFeature(
    ShapeReader('../shapefiles/china.shp').geometries(),
    proj,
    linewidth=0.3,
    edgecolor='grey',
    facecolor='none'
)


def plot_total_mean(df,obs_type):
    fig, axs = plot.subplots(ncols=1, nrows=1,figsize=(9,6),proj=proj)
    axs.add_feature(china)
    axs.format(
        lonlim=(args.min_lon, args.max_lon),
        latlim=(args.min_lat, args.max_lat),
        labels=True,
    )

    raob = df.query(f'obs_type=="{obs_type}"')
    
    raob_sfc = raob.groupby(['lon', 'lat']).agg({'obs_impact': ['mean']}).reset_index()
    raob_sfc.columns = ['lon', 'lat', 'mean_obs_impact']
    
    im = axs.scatter(raob_sfc['lon'], raob_sfc['lat'], marker='o', c=raob_sfc['mean_obs_impact'], s=10, cmap=cmaps.BlWhRe, levels=np.linspace(-10000, 10000, 21))
    axs.colorbar(im, loc='b')
    axs.format(
                title='EFSO {} column mean   {}'.format(obs_type,time_i.strftime('%Y-%m-%d %HZ'))
    )

    fig.savefig(f'{args.output}_{obs_type.lower()}_mean.png', dpi=200, bbox_inches='tight')

def plot_p_barh(df, obs_type):
    fig, axs = plot.subplots(figsize=(6,6))

    raob = df.query(f'obs_type=="{obs_type}"')

    raob_p = raob.groupby(pd.cut(raob['p'], (100,200,300,400,500,550,600,650,700,750,800,850,900,950,1000))).agg({'obs_impact': ['sum']}).reset_index()
    axs.barh(raob_p.index.values[::-1], raob_p['obs_impact'].values[::-1,0], align='center')
    axs.set_yticks(raob_p.index[::-1])
    axs.set_yticklabels(raob_p['p'].values[::-1])
    axs.format(title='EFSO {} {}'.format(obs_type,time_i.strftime('%Y-%m-%d %HZ')))
    fig.savefig(f'{args.output}_{obs_type.lower()}_p_barh.png', dpi=200, bbox_inches='tight')

def plot_var_p_barh(df, obs_type):
    if obs_type == 'RAOB':
        var = ['U','V','T','Q']
    elif obs_type == 'AMDAR':
        var = ['U','V','T']
    elif obs_type == 'PROFILER':
        var = ['U','V']
    else:
        return
    for ivar in var:
 
        fig, axs = plot.subplots(figsize=(6,6))

        raob = df[ (df['obs_type'] == obs_type) & (df['var'] == ivar )]

        raob_p = raob.groupby(pd.cut(raob['p'], (100,200,300,400,500,550,600,650,700,750,800,850,900,950,1000))).agg({'obs_impact': ['mean']}).reset_index()
        axs.barh(raob_p.index.values[::-1], raob_p['obs_impact'].values[::-1,0], align='center')
        axs.set_yticks(raob_p.index[::-1])
        axs.set_yticklabels(raob_p['p'].values[::-1])
        axs.format(title='EFSO {} {}  {}'.format(obs_type, ivar, time_i.strftime('%Y-%m-%d %HZ')))
        fig.savefig(f'{args.output}_{obs_type.lower()}_{ivar}_p_barh.png', dpi=200, bbox_inches='tight')


def plot_var_mean(df, obs_type):
    if obs_type == 'RAOB':
        var = ['U','V','T','Q']
    elif obs_type == 'AMDAR':
        var = ['U','V','T']
    elif obs_type == 'PROFILER':
        var = ['U','V']
    else:
        return
    for ivar in var:
        fig, axs = plot.subplots(ncols=1, nrows=1,figsize=(9,6),proj=proj)
        axs.add_feature(china)
        axs.format(
            lonlim=(args.min_lon, args.max_lon),
            latlim=(args.min_lat, args.max_lat),
            labels=True,
        )

        raob = df[ (df['obs_type'] == obs_type) & (df['var'] == ivar )]
        raob_sfc = raob.groupby(['lon', 'lat']).agg({'obs_impact': ['mean']}).reset_index()
        raob_sfc.columns = ['lon', 'lat', 'mean_obs_impact']
    
        im = axs.scatter(raob_sfc['lon'], raob_sfc['lat'], marker='o', c=raob_sfc['mean_obs_impact'], s=10, cmap=cmaps.BlWhRe, levels=np.linspace(-10000, 10000, 21))
        axs.colorbar(im, loc='b')
        axs.format(
                title='EFSO {} {} mean   {}'.format(obs_type, ivar, time_i.strftime('%Y-%m-%d %HZ'))
        )

        fig.savefig(args.output + '_{}_{}_mean.png'.format(obs_type.lower(), ivar), dpi=200, bbox_inches='tight')

def plot_diff_obs_type(df):
    fig, axs = plot.subplots(figsize=(8,6))

    raob_p = df.groupby( ['obs_type'] ).agg({'obs_impact': ['mean']}).reset_index()
    axs.barh(raob_p.index.values[::-1], raob_p['obs_impact'].values[::-1,0], height=0.2, align='center')
    axs.set_yticks(raob_p.index[::-1])
    axs.set_yticklabels(raob_p['obs_type'].values[::-1])
    axs.format(title='EFSO obs type {}'.format(time_i.strftime('%Y-%m-%d %HZ')))
    fig.savefig(f'{args.output}_obs_type_barh.png', dpi=200, bbox_inches='tight')


# -------------------------

efso_data = []
for line in open(args.input).readlines():
    cols = line.split()
    efso_data.append({
        'obs_type': cols[0],
        'var': cols[1],
        'sid': cols[2],
        'lon': float(cols[3]),
        'lat': float(cols[4]),
        'p': float(cols[5]),
        'obs_impact': float(cols[6]),
        'obs_value': float(cols[7]),
        'obs_inc': float(cols[8])
    })

df = pd.DataFrame(efso_data)

df = df[ (df['p'] >100) & (df['p'] <= 1000)]
print(df[df['obs_impact'] > 10000])
print(len(df[df['obs_impact'] > 10000]))

plot_total_mean(df,"RAOB")
plot_total_mean(df,"AMDAR")
plot_total_mean(df,"PROFILER")
plot_p_barh(df,"RAOB")
plot_p_barh(df,"AMDAR")
plot_p_barh(df,"PROFILER")
plot_var_p_barh(df,"RAOB")
plot_var_p_barh(df,"AMDAR")
plot_var_p_barh(df,"PROFILER")
plot_var_mean(df, "RAOB")
plot_var_mean(df, "AMDAR")
plot_var_mean(df, "PROFILER")
plot_diff_obs_type(df)
